package shardkv

// import "../shardmaster"
import (
	"bytes"
	"log"
	"sync"
	"time"

	"../labgob"
	"../labrpc"
	"../raft"
	"../shardmaster"
)

const debug = 1

//DPrintf :
func DPrintf(format string, a ...interface{}) (n int, err error) {
	if debug > 0 {
		log.Printf(format, a...)
	}
	return
}

//Op :
type Op struct {
	// Your definitions here.
	// Field names must start with capital letters,
	// otherwise RPC will break.
	ClientID  int64
	ReqNum    int
	Operation string
	Key       string
	Value     string

	NewConfig     shardmaster.Config
	ShardTransfer ShardMigration
}

//Configuration :
type Configuration struct {
	NewConfig shardmaster.Config
}

//ShardMigration :
type ShardMigration struct {
	ConfigNum  int
	Shard      int
	Operation  string
	Data       map[string]string
	ReqHistory map[int64]int
}

//Result :
type Result struct {
	Err   Err
	value string
}

//ShardKV :
type ShardKV struct {
	mu           sync.Mutex
	me           int
	rf           *raft.Raft
	applyCh      chan raft.ApplyMsg
	make_end     func(string) *labrpc.ClientEnd
	gid          int
	masters      []*labrpc.ClientEnd
	maxraftstate int // snapshot if log grows this big

	// Your definitions here.
	sm            *shardmaster.Clerk
	config        shardmaster.Config
	database      []map[string]string
	reqHistory    []map[int64]int //clientid -> ReqNum
	shardTransfer map[int]bool    //wait for shard transfering
	shardHistory  map[int]int     //shard -> config.Num
	lastconfignum int

	requestCh chan Result
	shardCh   map[int]chan TransferReply
}

func (kv *ShardKV) queryReqHistory(clientid int64, curNum int, shard int) bool {
	maxNum, ok := kv.reqHistory[shard][clientid]
	if !ok {
		kv.reqHistory[shard][clientid] = -1
		return false
	}
	return maxNum >= curNum

}

//Get :
func (kv *ShardKV) Get(args *GetArgs, reply *GetReply) {
	// Your code here.
	kv.mu.Lock()
	shard := key2shard(args.Key)
	ok := kv.queryReqHistory(args.ClientID, args.ReqNum, shard)
	if kv.config.Shards[shard] != kv.gid || len(kv.shardTransfer) > 0 { //not responsible for this shard or wait for shard transfering
		DPrintf("[%v, %v]: get wronggroup len(kv.shardTransfer) %v", kv.gid, kv.me, len(kv.shardTransfer))
		reply.Err = ErrWrongGroup
		kv.mu.Unlock()
		return
	}
	kv.mu.Unlock()
	command := Op{args.ClientID, args.ReqNum, "Get", args.Key, "", shardmaster.Config{}, ShardMigration{}}
	if !ok {
		term1, _, isleader := kv.rf.Start(command)
		if isleader {
			//kv.mu.Unlock()
			var result Result
			select {
			case result = <-kv.requestCh:
			case <-time.After(150 * time.Millisecond):
				//log.Printf("kv[%v]: get time out\n", kv.me)
				reply.Err = ErrWrongLeader
				return
			}
			if term2, _ := kv.rf.GetState(); term2 != term1 {
				reply.Err = ErrWrongLeader
				return
			}
			if result.Err != ErrWrongGroup {
				reply.Value = result.value
				if reply.Value == "" {
					reply.Err = ErrNoKey
				} else {
					reply.Err = OK
				}
			} else {
				reply.Err = ErrWrongGroup
			}
			DPrintf("[%v, %v]: Get Success value %v, shard %v, database %v", kv.gid, kv.me, reply.Value, shard, kv.database)
		} else {
			reply.Err = ErrWrongLeader
		}
	} else {
		kv.mu.Lock()
		result := kv.database[shard][args.Key]
		kv.mu.Unlock()
		if result == "" {
			reply.Err = ErrNoKey
		} else {
			reply.Err = OK
		}
		reply.Value = result
		DPrintf("[%v, %v]: Get Success2 value %v, shard %v, database %v", kv.gid, kv.me, reply.Value, shard, kv.database)
	}
}

//PutAppend :
func (kv *ShardKV) PutAppend(args *PutAppendArgs, reply *PutAppendReply) {
	// Your code here.
	kv.mu.Lock()
	shard := key2shard(args.Key)
	ok := kv.queryReqHistory(args.ClientID, args.ReqNum, shard)
	if kv.config.Shards[shard] != kv.gid || len(kv.shardTransfer) > 0 { //not responsible for this shard or wait for shard transfering
		DPrintf("[%v, %v]: PutAppend wronggroup len(kv.shardTransfer) %v", kv.gid, kv.me, len(kv.shardTransfer))
		reply.Err = ErrWrongGroup
		kv.mu.Unlock()
		return
	}
	kv.mu.Unlock()
	command := Op{args.ClientID, args.ReqNum, args.Op, args.Key, args.Value, shardmaster.Config{}, ShardMigration{}}
	if !ok {
		term1, _, isleader := kv.rf.Start(command)
		if isleader {
			var result Result
			//DPrintf("[%v, %v]: start Put key %v, value %v", kv.gid, kv.me, args.Key, args.Value)
			select {
			case result = <-kv.requestCh:
			case <-time.After(150 * time.Millisecond):
				//log.Printf("kv[%v]: putappend time out", kv.me)
				reply.Err = ErrWrongLeader
				return
			}
			if term2, _ := kv.rf.GetState(); term2 != term1 {
				reply.Err = ErrWrongLeader
				return
			}
			reply.Err = result.Err
		} else {
			reply.Err = ErrWrongLeader
		}
	} else {
		reply.Err = OK
	}
}

//TransferShard :
func (kv *ShardKV) TransferShard(args *TransferArgs, reply *TransferReply) {
	kv.mu.Lock()
	historyConfigNum := kv.shardHistory[args.Shard]
	lastconfignum := kv.lastconfignum
	kv.mu.Unlock()
	if historyConfigNum != args.CurConfigNum-1 || lastconfignum < args.CurConfigNum {
		DPrintf("[%v, %v]: historyConfigNum %v != args.CurConfigNum-1 %v , args.Shard %v", kv.gid, kv.me, historyConfigNum, args.CurConfigNum-1, args.Shard)
		reply.Success = false
		return
	}
	command := Op{}
	command.Operation = "Transfer"
	command.ShardTransfer = ShardMigration{historyConfigNum, args.Shard, "Send", nil, nil}
	_, _, isleader := kv.rf.Start(command)
	if isleader {
		//kv.mu.Unlock()
		var result TransferReply
		select {
		case result = <-kv.shardCh[args.Shard]:
		case <-time.After(150 * time.Millisecond):
			log.Printf("kv[%v, %v]: get time out\n", kv.gid, kv.me)
			reply.Success = false
			return
		}
		reply.Data = result.Data
		reply.ReqHistory = result.ReqHistory
		reply.ShardConfigNum = result.ShardConfigNum
		reply.Success = true
		DPrintf("[%v, %v]: send shard data %v, database %v", kv.gid, kv.me, args.Shard, reply.Data)
	} else {
		reply.Success = false
	}

}

//
func (kv *ShardKV) getConfig() {
	for {
		time.Sleep(90 * time.Millisecond)
		_, isleader := kv.rf.GetState()
		if !isleader {
			continue
		}
		kv.mu.Lock()
		if len(kv.shardTransfer) > 0 { //Process re-configurations one at a time, in order
			kv.mu.Unlock()
			continue
		}
		newconfig := kv.sm.Query(kv.lastconfignum + 1)
		if newconfig.Num != kv.lastconfignum+1 {
			kv.mu.Unlock()
			continue
		}
		kv.mu.Unlock()
		DPrintf("[%v, %v]: config changed! config %v", kv.gid, kv.me, newconfig)
		command := Op{}
		//command := Configuration{}
		command.Operation = "ConfigChanged"
		command.NewConfig = newconfig
		kv.rf.Start(command)
	}
}

func (kv *ShardKV) startShardTransfer(newconfig shardmaster.Config, oldconfig shardmaster.Config) {
	if newconfig.Num != 1 {
		for i := 0; i < len(newconfig.Shards); i++ {
			if oldconfig.Shards[i] != newconfig.Shards[i] { //shard owner changed!
				if oldconfig.Shards[i] != kv.gid && newconfig.Shards[i] == kv.gid && kv.shardHistory[i] < newconfig.Num { //become shard's new owner and haven't recv shard data
					//DPrintf("[%v, %v]: kv.shardHistory[i] %v, newconfig.Num %v", kv.gid, kv.me, kv.shardHistory[i], newconfig.Num)
					args := TransferArgs{newconfig.Num, kv.shardHistory[i], i}
					servers := oldconfig.Groups[oldconfig.Shards[i]] //shard of old owner
					kv.shardTransfer[i] = false
					go kv.pullShard(args, servers, newconfig.Num)
				}
			} else if newconfig.Shards[i] == kv.gid {
				kv.shardHistory[i] = newconfig.Num
			}
		}
	} else {
		for i := 0; i < len(newconfig.Shards); i++ {
			if newconfig.Shards[i] == kv.gid {
				kv.shardHistory[i] = 1 //init
			}
		}
	}
}

func (kv *ShardKV) pullShard(args TransferArgs, servers []string, newconfignum int) {
	for {
		if _, isleader := kv.rf.GetState(); !isleader {
			return
		}
		for si := 0; si < len(servers); si++ {
			reply := TransferReply{}
			srv := kv.make_end(servers[si])
			ok := srv.Call("ShardKV.TransferShard", &args, &reply)
			//DPrintf("[%v, %v]: ok %v, reply.Success %v ", kv.gid, kv.me, ok, reply.Success)
			if ok && reply.Success {
				kv.mu.Lock()
				if reply.ShardConfigNum > kv.shardHistory[args.Shard] {
					kv.mu.Unlock()
					DPrintf("[%v, %v]: recv TransferShard reply %v", kv.gid, kv.me, reply)
					command := Op{}
					command.Operation = "Transfer"
					command.ShardTransfer.Operation = "Recv"
					command.ShardTransfer.Data = make(map[string]string)
					command.ShardTransfer.ReqHistory = make(map[int64]int)
					for k, v := range reply.Data {
						command.ShardTransfer.Data[k] = v
					}
					for k, v := range reply.ReqHistory {
						command.ShardTransfer.ReqHistory[k] = v
					}
					command.ShardTransfer.ConfigNum = newconfignum
					command.ShardTransfer.Shard = args.Shard
					if _, _, isleader2 := kv.rf.Start(command); isleader2 {
						//DPrintf("[%v, %v]: Start command success", kv.gid, kv.me)
						return
					}
					//DPrintf("[%v, %v]: I am not leader recv TransferShard reply %v", kv.gid, kv.me, reply)
				} else {
					DPrintf("[%v, %v]: old shard data %v , reply.ShardConfigNum %v <= kv.shardHistory[args.Shard] %v", kv.gid, kv.me, args.Shard, reply.ShardConfigNum, kv.shardHistory[args.Shard])
					kv.mu.Unlock()
				}
			}
		}
		kv.mu.Lock()
		if _, ok := kv.shardTransfer[args.Shard]; !ok {
			kv.mu.Unlock()
			return
		}
		kv.mu.Unlock()
		time.Sleep(50 * time.Millisecond)
	}
}

func (kv *ShardKV) checkShardTransfer() {
	for {
		time.Sleep(200 * time.Millisecond)
		if _, isleader := kv.rf.GetState(); !isleader {
			continue
		}
		kv.mu.Lock()
		if kv.lastconfignum <= 1 {
			kv.mu.Unlock()
			continue
		}
		oldconfig := kv.sm.Query(kv.lastconfignum - 1)
		for k := range kv.shardTransfer {
			args := TransferArgs{kv.config.Num, kv.shardHistory[k], k}
			servers := oldconfig.Groups[oldconfig.Shards[k]] //shard of old owner
			go kv.pullShard(args, servers, kv.config.Num)
		}
		kv.mu.Unlock()
	}

}

//applyRoutine :
func (kv *ShardKV) applyRoutine() {
	for {
		applymsg := <-kv.applyCh
		if applymsg.CommandValid {
			kv.mu.Lock()
			//var op Op
			//DPrintf("[%v]: command %v", kv.me, applymsg.Command)
			if applymsg.CommandIndex == 0 {
				kv.mu.Unlock()
				continue
			}
			op, isOp := applymsg.Command.(Op)
			DPrintf("[%v, %v]: index %v, command %v", kv.gid, kv.me, applymsg.CommandIndex, applymsg.Command)
			//if _, ok := kv.stateMachine[op.ID]; ok || !isOp { //duplicate
			if isOp {
				shard := key2shard(op.Key)
				if ok := kv.queryReqHistory(op.ClientID, op.ReqNum, shard); ok {
					kv.mu.Unlock()
					continue
				}
				//DPrintf("[%v, %v]: command %v", kv.gid, kv.me, applymsg.Command)
				if op.Operation == "ConfigChanged" {
					if op.NewConfig.Num == kv.lastconfignum+1 { //in case double config check
						kv.startShardTransfer(op.NewConfig, kv.config)
						kv.config = op.NewConfig
						kv.lastconfignum = op.NewConfig.Num
					}
					//DPrintf("[%v, %v]: after startShardTransfer", kv.gid, kv.me)
					kv.mu.Unlock()
				} else if op.Operation == "Transfer" {
					if op.ShardTransfer.Operation == "Send" {
						var result TransferReply
						result.Data = make(map[string]string)
						result.ReqHistory = make(map[int64]int)
						for k, v := range kv.database[op.ShardTransfer.Shard] {
							result.Data[k] = v
						}
						for k, v := range kv.reqHistory[op.ShardTransfer.Shard] {
							result.ReqHistory[k] = v
						}
						result.ShardConfigNum = kv.shardHistory[op.ShardTransfer.Shard]
						kv.mu.Unlock()
						select {
						case kv.shardCh[op.ShardTransfer.Shard] <- result: //try to send channel to putappend or get RPC
						default:
						}
					} else if _, ok := kv.shardTransfer[op.ShardTransfer.Shard]; ok {
						kv.database[op.ShardTransfer.Shard] = make(map[string]string)
						kv.reqHistory[op.ShardTransfer.Shard] = make(map[int64]int)
						for k, v := range op.ShardTransfer.Data {
							kv.database[op.ShardTransfer.Shard][k] = v
						}
						for k, v := range op.ShardTransfer.ReqHistory {
							kv.reqHistory[op.ShardTransfer.Shard][k] = v
						}
						//kv.database[op.ShardTransfer.Shard] = op.ShardTransfer.Data
						kv.shardHistory[op.ShardTransfer.Shard] = op.ShardTransfer.ConfigNum
						delete(kv.shardTransfer, op.ShardTransfer.Shard)
						DPrintf("[%v, %v]: recv shard data %v", kv.gid, kv.me, op.ShardTransfer.Shard)
						kv.mu.Unlock()
					} else {
						DPrintf("[%v, %v]: old shard data %v", kv.gid, kv.me, op.ShardTransfer.Shard)
						kv.mu.Unlock()
					}
				} else {
					var result Result
					shard := key2shard(op.Key)
					if kv.config.Shards[shard] != kv.gid || len(kv.shardTransfer) > 0 {
						result.Err = ErrWrongGroup
					} else {
						switch op.Operation {
						case "Put":
							kv.database[shard][op.Key] = op.Value
							result.Err = OK
							//DPrintf("[%v]: Put finish", kv.me)
						case "Append":
							tmp := kv.database[shard][op.Key]
							kv.database[shard][op.Key] = tmp + op.Value
							result.Err = OK
							//DPrintf("[%v]: Append finish", kv.me)
						case "Get":
							//DPrintf("[%v]: Get finish", kv.me)
							value, ok := kv.database[shard][op.Key]
							if ok {
								result.value = value
							} else {
								result.value = ""
							}
						default:
						}
						kv.reqHistory[shard][op.ClientID] = op.ReqNum
					}
					kv.mu.Unlock()
					select {
					case kv.requestCh <- result: //try to send channel to putappend or get RPC
					default:
					}
				}
			}
		} else if applymsg.IsSnapshot { //recv snapshot
			kv.readSnapshot(applymsg.Snapshots)
			DPrintf("kv[%v, %v]: recv snapshot database %v\n", kv.gid, kv.me, kv.database)
		} else {
			kv.mu.Lock()
			snapshotmsg := raft.ApplyMsg{}
			if kv.maxraftstate != -1 && applymsg.RaftStateSize > kv.maxraftstate { //send snapshot
				DPrintf("[%v, %v]: before createsnapshot", kv.gid, kv.me)
				writesnapshot := new(bytes.Buffer)
				enc := labgob.NewEncoder(writesnapshot)
				enc.Encode(kv.config)
				enc.Encode(kv.database)
				enc.Encode(kv.reqHistory)
				enc.Encode(kv.shardTransfer)
				enc.Encode(kv.shardHistory)
				enc.Encode(kv.lastconfignum)
				snapshot := writesnapshot.Bytes()
				snapshotmsg.CommandValid = true
				snapshotmsg.Snapshots = snapshot
				DPrintf("[%v, %v]: after createsnapshot database %v", kv.gid, kv.me, kv.database)
			} else {
				snapshotmsg.CommandValid = false
			}
			kv.mu.Unlock()
			kv.applyCh <- snapshotmsg //communicate with raft to confirm whether create snapshot or not and snapshot index
		}
	}
}

func (kv *ShardKV) readSnapshot(data []byte) {
	kv.mu.Lock()
	defer kv.mu.Unlock()
	r := bytes.NewBuffer(data)
	d := labgob.NewDecoder(r)
	var config shardmaster.Config
	var db []map[string]string
	var reqhistory []map[int64]int
	var shtran map[int]bool
	var sh map[int]int
	var num int
	if d.Decode(&config) != nil ||
		d.Decode(&db) != nil ||
		d.Decode(&reqhistory) != nil ||
		d.Decode(&shtran) != nil ||
		d.Decode(&sh) != nil ||
		d.Decode(&num) != nil {
		DPrintf("kv[%v, %v]: readSnapshot error!", kv.gid, kv.me)
	} else {
		kv.config = config
		kv.database = db
		kv.reqHistory = reqhistory
		kv.shardTransfer = shtran
		kv.shardHistory = sh
		kv.lastconfignum = num
	}
}

//Kill :
// the tester calls Kill() when a ShardKV instance won't
// be needed again. you are not required to do anything
// in Kill(), but it might be convenient to (for example)
// turn off debug output from this instance.
//
func (kv *ShardKV) Kill() {
	kv.rf.Kill()
	// Your code here, if desired.
}

//StartServer :
// servers[] contains the ports of the servers in this group.
//
// me is the index of the current server in servers[].
//
// the k/v server should store snapshots through the underlying Raft
// implementation, which should call persister.SaveStateAndSnapshot() to
// atomically save the Raft state along with the snapshot.
//
// the k/v server should snapshot when Raft's saved state exceeds
// maxraftstate bytes, in order to allow Raft to garbage-collect its
// log. if maxraftstate is -1, you don't need to snapshot.
//
// gid is this group's GID, for interacting with the shardmaster.
//
// pass masters[] to shardmaster.MakeClerk() so you can send
// RPCs to the shardmaster.
//
// make_end(servername) turns a server name from a
// Config.Groups[gid][i] into a labrpc.ClientEnd on which you can
// send RPCs. You'll need this to send RPCs to other groups.
//
// look at client.go for examples of how to use masters[]
// and make_end() to send RPCs to the group owning a specific shard.
//
// StartServer() must return quickly, so it should start goroutines
// for any long-running work.
//
func StartServer(servers []*labrpc.ClientEnd, me int, persister *raft.Persister, maxraftstate int, gid int, masters []*labrpc.ClientEnd, make_end func(string) *labrpc.ClientEnd) *ShardKV {
	// call labgob.Register on structures you want
	// Go's RPC library to marshall/unmarshall.
	labgob.Register(Op{})

	kv := new(ShardKV)
	kv.me = me
	kv.maxraftstate = maxraftstate
	kv.make_end = make_end
	kv.gid = gid
	kv.masters = masters

	// Your initialization code here.

	// Use something like this to talk to the shardmaster:
	// kv.mck = shardmaster.MakeClerk(kv.masters)

	kv.applyCh = make(chan raft.ApplyMsg)
	kv.rf = raft.Make(servers, me, persister, kv.applyCh)

	kv.sm = shardmaster.MakeClerk(kv.masters)
	kv.database = make([]map[string]string, 10)
	kv.reqHistory = make([]map[int64]int, 10)
	kv.shardTransfer = make(map[int]bool)
	kv.shardHistory = make(map[int]int)
	for i := 0; i < 10; i++ {
		kv.database[i] = make(map[string]string)
		kv.reqHistory[i] = make(map[int64]int)
	}
	kv.lastconfignum = 0
	kv.requestCh = make(chan Result)
	kv.shardCh = make(map[int]chan TransferReply)
	for i := 0; i < 10; i++ {
		kv.shardCh[i] = make(chan TransferReply)
	}
	kv.readSnapshot(persister.ReadSnapshot())
	//kv.config = kv.sm.Query(-1)
	go kv.getConfig()
	go kv.applyRoutine()
	go kv.checkShardTransfer()

	return kv
}
